#!/usr/bin/env python3
import os
import torch
import rclpy
import math
import numpy as np

from rclpy.node import Node
from loguru import logger
from geometry_msgs.msg import PoseStamped
from mavros_msgs.msg import State, AttitudeTarget, RCIn
from nav_msgs.msg import Odometry
from sensor_msgs.msg import Imu # type sensor_msgs/msg/Imu

from std_msgs.msg import String,UInt32

from .lib.rltrain import attitudePolicy
from scipy.spatial.transform import Rotation
   
class AttitudeNode(Node):
    # Transformation matrix from px4_body to ros_body    
    matrix_px4body2rosbody = np.array([[1, 0, 0],
                                        [0, -1, 0], 
                                        [0, 0, -1]])

    # Transformation matrix from ros_local to px4_local  
    matrix_roslocal2px4local = np.array([[0, 1, 0],
                                        [1, 0, 0], 
                                        [0, 0, -1]])

    thrust = 0.0
    n3xd = 0.0
    n3yd = 0.0
    n3zd = 0.0
    def __init__(self, name):
        super().__init__(name)
        self.get_logger().info(f"{name} started")
        
        # 订阅话题 /mavros/local_position/pose
        # self.attitude_sub = self.create_subscription(PoseStamped, "/mavros/local_position/pose", self.state_callback, rclpy.qos.qos_profile_sensor_data)

        # 订阅话题 /mavros/imu/data
        self.imu_sub = self.create_subscription(Imu, "/mavros/imu/data", self.imu_callback, rclpy.qos.qos_profile_sensor_data)

        # 订阅话题 /mavros/rc/in
        self.rcin_sub = self.create_subscription(RCIn, "/mavros/rc/in", self.rcin_callback, rclpy.qos.qos_profile_sensor_data)

        # 发布话题
        self.attitude_target_pub = self.create_publisher(AttitudeTarget,"/mavros/setpoint_raw/attitude", rclpy.qos.qos_profile_sensor_data)

        # 创建定时循环
        # self.timer = self.create_timer(1, self.timer_callback)
        
        # 加载模型
        model_path = os.path.expanduser('~/attitude_net_dict.pt')
        device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
        self.attPolicy = attitudePolicy().to(device)
        try:
            self.attPolicy.load_state_dict(torch.load(model_path, map_location=device))
            self.get_logger().warn(f"Model loaded successfully on device: {device}")
        except FileNotFoundError:
            raise FileNotFoundError(f"Model file not found: {model_path}")
        except Exception as e:
            raise RuntimeError(f"Failed to load model from {model_path}: {str(e)}")

        self.attPolicy.eval()

    def imu_callback(self, msg):
        # 获取PX4坐标系下的旋转矩阵
        R_eb_px4 = self.rotation_ros2px4(np.array([msg.orientation.x,msg.orientation.y,msg.orientation.z,msg.orientation.w]))
        R_be = R_eb_px4.as_matrix().T
        
        # 获取PX4坐标系下的角速度
        rate_px4 = self.matrix_px4body2rosbody.T @ np.array([[msg.angular_velocity.x],[msg.angular_velocity.y],[msg.angular_velocity.z]])

        # 输入带入神经网络计算其输出
        self.n3zd = math.sqrt(1 - self.n3xd**2 - self.n3yd**2)
        nn_state = np.array([msg.angular_velocity.z, R_be[0,2], R_be[1,2], R_be[2,2], self.n3xd, self.n3yd, self.n3zd])
        nn_state = torch.from_numpy(nn_state).to(torch.float32).to(device='cuda:0')
        nn_output = self.attPolicy(nn_state)
        omega_2 = nn_output.cpu().detach().numpy()
        # Ensure omega_2 is 2D and has correct shape
        omega_2 = omega_2.reshape(1, -1) if omega_2.ndim == 1 else omega_2
        # self.get_logger().info(f"nn_output is {omega_2[0][0]}")

        # send ctrl cmd to px4
        att_tgt_msg = AttitudeTarget()
        att_tgt_msg.header.stamp = self.get_clock().now().to_msg()
        att_tgt_msg.header.frame_id = "attitude_node"
        att_tgt_msg.type_mask = AttitudeTarget.IGNORE_ATTITUDE
        att_tgt_msg.body_rate.x = float(omega_2[0,0]) if omega_2.shape[1] > 0 else 0.0
        att_tgt_msg.body_rate.y = -float(omega_2[0,1]) if omega_2.shape[1] > 1 else 0.0
        att_tgt_msg.body_rate.z = 0.0
        att_tgt_msg.thrust = self.thrust
        self.attitude_target_pub.publish(att_tgt_msg)
    
    def rcin_callback(self, msg):
        self.thrust = (msg.channels[2] - 1500.0)*0.001 + 0.5 #不同的遥控器油门通道可能存在反相，需要根据实际情况调整
        self.n3xd = -(msg.channels[1] - 1500.0)*0.001 * 0.2
        self.n3yd = (msg.channels[0] - 1500.0)*0.001 * 0.2
        # self.get_logger().info(f"n3xd is {self.n3xd}, n3yd is {self.n3yd}")

    # def timer_callback(self):
    #     test_state = np.array([0, 0, 0, 0, 0, 0, 0])
    #     test_state = torch.from_numpy(test_state).to(torch.float32).to(device='cuda:0')
    #     test_output = self.attPolicy(test_state)
    #     logger.info(f"test_state is {test_state.cpu().numpy()}")
    #     logger.info(f"test_output is {test_output.cpu().detach().numpy()[0,:]}")

    def rotation_ros2px4(self, quat_from_mavros)-> Rotation:
        """
        quat_from_mavros: quaternion [x, y, z, w] in ROS coordinate system
        """
        # Rotation from ros_body to ros_local
        R_eb_ros = Rotation.from_quat(quat_from_mavros) 

        # Transformation matrix from px4_body to px4_local
        matrix_eb_px4 = self.matrix_roslocal2px4local @ R_eb_ros.as_matrix() @ self.matrix_px4body2rosbody 

        R_eb_px4 = Rotation.from_matrix(matrix_eb_px4)
        return R_eb_px4

def main(args=None):
    rclpy.init(args=args)
    node = AttitudeNode("attitude_node")
    node.get_logger().info("attitude node started")
    rclpy.spin(node)
    rclpy.shutdown()

